#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
 SO3 base type - implements SO3 class but is storage agnostic.

 SO(3) is the group of rotations in 3d. As a matrix group, it is the set of
 matrices which are orthogonal such that ``R * R' = I`` (with ``R'`` being the
 transpose of ``R``) and have a positive determinant. In particular, the
 determinant is 1. Internally, the group is represented as a unit quaternion.
 Unit quaternion can be seen as members of the special unitary group SU(2).
 SU(2) is a double cover of SO(3). Hence, for every rotation matrix ``R``,
 there exists two unit quaternion: ``(r, v)`` and ``(r, -v)``, with ``r`` the
 real part and ``v`` being the imaginary 3-vector part of the quaternion.

 SO(3) is a compact, but non-commutative group. First it is compact since the
 set of rotation matrices is a closed and bounded set. Second it is
 non-commutative since the equation ``R_1 * R_2 = R_2 * R_1`` does not hold in
 general. For example rotating an object by some degrees about its ``x``-axis
 and then by some degrees about its y axis, does not lead to the same
 orienation when rotation first about ``y`` and then about ``x``.

 Class invairant: The 2-norm of ``unit_quaternion`` must be close to 1.
 Technically speaking, it must hold that:

   ``|unit_quaternion().squaredNorm() - 1| <= Constants<Scalar>::epsilon()``.
'''

import numpy as np
import eulerangles as euler
import Quaternion as quat
import transform as tf

class SO3_base():

    def __init__(self):
        self.rotation = np.eye(3)

    def __add__(self, other):
        self.rotation.dot(other.rotation)

    def __sub__(self, other):
        self.rotation.dot(np.linalg.inv(other.rotation))

    def hat(self, v):
        """
        给定向量，输出其hat矩阵
        :param
        v: 3d vector
        :return
        out: 3d skew matrix of v
        @exsamples
        the input vector is:
            [ 0.00912655  0.02745836  0.07600277]
        the output vector is:
            [[ 0.         -0.07600277  0.02745836]
             [ 0.07600277  0.         -0.00912655]
             [-0.02745836  0.00912655  0.        ]]
        """
        if len(v) > 4 or len(v) < 0:
            raise LookupError('please input proper vector')
        if len(v) == 4:
            v = v[:3] / v[3]
        skv = np.roll(np.roll(np.diag(v.flatten()), 1, 1), -1, 0)
        return skv - skv.T

    def exp(self, so3):
        """
        :param
        so3: 3d vector
        :return
        exp(so3^): 3d SO3 map matrix
        @exsamples: BCH formula
        input so3 vector:
            [ 0.00912655  0.02745836  0.07600277]
        output SO3 matrix:
            [[ 0.99673661 -0.07579379  0.02777473]
             [ 0.07604425  0.99707176 -0.00807362]
             [-0.02708147  0.01015938  0.9995816 ]]
        """
        phi = so3
        theta = np.linalg.norm(phi)
        if theta == 0:
            SO3 = np.eye(3)
        else:
            alpha = phi / theta
            alpha = np.nan_to_num(alpha)
            alpha = np.expand_dims(alpha, axis=0)  # row vector
            SO3 = np.cos(theta) * np.eye(3) + (1 - np.cos(theta)) * alpha.T.dot(alpha) + np.sin(theta) * self.hat(alpha)
        return SO3

    def log(self, SO3):
        '''
        :param
        SO3: 3d SO3 matrix
        :return
        so3: 3d so3 vector
        @exsamples
        input matrix:
            [[ 0.99673661 -0.07579379  0.02777473]
             [ 0.07604425  0.99707176 -0.00807362]
             [-0.02708147  0.01015938  0.9995816 ]]
        output vector:
            [ 0.00912655  0.02745836  0.07600277]
        '''
        rot = SO3[:3, :3]
        theta = np.arccos((np.trace(rot) - 1) / 2)
        lnR = theta / 2 / np.sin(theta) * (rot - rot.T)
        alpha = np.array([-lnR[1, 2], lnR[0, 2], -lnR[0, 1]])
        so3 = alpha
        return so3

    def inv(self, SO3):
        '''
        :param
        SO3: rotation matrix
        :return
        out: the inverse of SO3
        '''
        inv = SO3.T
        return inv

    def rotation_matrix(self):
        return self.exp()

    def set_euler(self, _euler):
        x, y, z = _euler
        self.rotation = euler.euler2mat(z, x, y)

    def set_quaternion(self, _quaternion):
        self.rotation = quat.quat2mat(_quaternion)

    def set_axis_angle(self, theta, axis):
        mat = tf.rotation_matrix(theta, axis)
        self.rotation = mat[:3, :3]